# 将训练过程中的关键指标可视化并保存为图片

import pandas as pd
import matplotlib.pyplot as plt
import os

# 1️⃣ 路径设置
csv_path = "logs/train/training_metrics.csv"
output_dir = "logs/train"

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

# 2️⃣ 读取 CSV 文件
df = pd.read_csv(csv_path)

# 检查有 timesteps 列
if 'timesteps' not in df.columns:
    raise ValueError("CSV 文件中必须包含 'timesteps' 列。")

# 3️⃣ 遍历除 timesteps 外的每一列
for col in df.columns:
    if col == 'timesteps':
        continue
    
    plt.figure(figsize=(10, 5))
    plt.plot(df['timesteps'], df[col], label=col, linewidth=1.8)
    plt.xlabel("Timesteps")
    plt.ylabel(col)
    plt.title(f"{col} over Timesteps")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    
    # 保存图片
    save_path = os.path.join(output_dir, f"{col}.png")
    plt.savefig(save_path)
    plt.close()
    print(f"[Saved] {save_path}")

print("✅ 所有图表已生成并保存至 logs/train/")
